from collections import namedtuple, Counter
from pathlib import Path
from typing import Union, Sequence

import PIL
import cv2
import mmengine
import numpy as np
from PIL import ImageFont

from flask_moudle.service.models import yolov8_seg, PIDNet
from utils.download import download_file

Cls = namedtuple('cls', ['name', 'id', 'color'])
Clss = [
    Cls('background', 0, (0, 0, 0)),
    Cls('intact', 1, (0, 255, 0)),
    Cls('slight', 2, (255, 255, 0)),
    Cls('severe', 3, (255, 140, 0)),
    Cls('collapse', 4, (255, 0, 0))
]
default_classes = [
    'intact',
    'slight',
    'severe',
    'collapse']
default_colors = [(0, 0, 0), (0, 255, 0), (255, 255, 0), (255, 140, 0), (255, 0, 0)]


def draw_bbox_label(image, boxes, labels, scores, show_bbox: bool = False, show_labels=False, show_conf=False,
                    colors=None, font_size=20, width=4):
    """
    在图像上绘制边界框并添加文字

    :param image: 输入图像
    :param boxes: 边界框
    :param labels: 标签
    :param scores: 得分
    :param show_bbox: 是否显示bbox
    :param show_labels: 是否显示labels
    :param show_conf: 是否显示置信度
    :param colors: 颜色数组
    :param font_size: 字体大小
    :param width: 线宽
    """
    # 创建 ImageDraw 对象

    if colors is None:
        colors = default_colors
    draw = ImageDraw.Draw(image)

    # 绘制边界框
    for bbox, label, score, color in zip(boxes, labels, scores, colors):
        draw.rectangle((bbox[0], bbox[1], (bbox[2]), (bbox[3])), outline=color, width=width)
        # 添加文字
        # 可以设置自动检测字体是否存在
        font = ImageFont.truetype('STXINWEI.TTF', font_size)

        text = ''
        if show_bbox is True:
            draw.textbbox((bbox[0], bbox[1]), text, font=font, align="center")
        if show_labels is True:
            text = f'{label}'
        if show_conf is True:
            text = f'{score * 100:.2f}'
        if show_conf and show_labels is True:
            text = f'{label:}{score * 100:.2f}'
        if show_conf and show_labels is False:
            return
        draw.text((bbox[0] + 20, bbox[1]), text, fill=color, font=font)


def draw_RGB_mask(mask: np.array([[]]), colors: [] = None) -> np.array([[[]]]):
    """

    :param mask: 掩码数组 (h,w)
    :param colors: 颜色数组
    :return: RGB融合图 (h,w,3)
    """
    if colors is None:
        colors = colors
    h, w = np.shape(mask)
    masks_RGB = np.reshape(np.array(colors, np.uint8)[np.reshape(mask, [-1])], [h, w, -1])
    return masks_RGB


from PIL import Image, ImageDraw


def get_putpalette(Clss, color_other=None):
    '''
    灰度图转8bit彩色图
    :param Clss:颜色映射表
    :param color_other:其余颜色设置
    :return:
    '''
    if color_other is None:
        color_other = [0, 0, 0]
    putpalette = []
    for cls in Clss:
        putpalette += list(cls.color)
    putpalette += color_other * (255 - len(Clss))
    return putpalette


def instance_seg_predict(source: PIL.Image, model='yolov8-seg', conf=0.25, iou=0.5, blend: float = 0.5, imgsz=None,
                         classes=None, show_bbox: bool = False, show_labels=False, show_conf=False) -> tuple:
    """

    :param source: 图片源
    :param model: 使用的模型
    :param conf: 置信度
    :param iou: iou
    :param blend: 图像混合值
    :param imgsz: 图像大小
    :param classes: 类别
    :param show_bbox: 显示bbox
    :param show_labels: 显示类别标签
    :param show_conf: 显示得分
    :return: 混合图，掩码图，分析结果
    """
    # Load  model
    if imgsz is None:
        imgsz = [1024, 1024]
    if classes is None:
        classes = default_classes
    origin_size = source.size
    if model == 'yolov8-seg':
        model_path='weights/YOLO-Seg/yolo-v8-seg-n.onnx'
        download_file('https://file.zouran.top/weights/YOLO-Seg/yolo-v8-seg-n.onnx',model_path)
        isExist, labels, scores, boxes, masks = yolov8_seg(source, conf=conf, iou=iou, classes=classes, imgsz=imgsz,model_path=model_path)
    elif model=='yolov9-seg':
        model_path='weights/YOLO-Seg/yolo-v9-seg-c.onnx'
        download_file('https://file.zouran.top/weights/YOLO-Seg/yolo-v9-seg-c.onnx',model_path)
        isExist, labels, scores, boxes, masks = yolov8_seg(source, conf=conf, iou=iou, classes=classes, imgsz=imgsz,model_path=model_path)
    else:
        model_path='weights/YOLO-Seg/yolo-v8-seg-n.onnx'
        download_file('https://file.zouran.top/weights/YOLO-Seg/yolo-v8-seg-n.onnx',model_path)
        isExist, labels, scores, boxes, masks = yolov8_seg(source, conf=conf, iou=iou, classes=classes, imgsz=imgsz,model_path=model_path)
    # 不存在则直接返回原图
    if isExist is False:
        return source, np.zeros((origin_size[1], origin_size[0]), dtype=np.uint8), None
    image = Image.new('L', origin_size, color=0)
    draw = ImageDraw.Draw(image)
    for mask, label in zip(masks, labels):
        draw.polygon(mask.flatten().tolist(), fill=label + 1)
        image.putpalette(get_putpalette(Clss))
    # save image
    mask_np = np.array(image)
    masks_RGB = draw_RGB_mask(np.array(image), default_colors)
    masks_ = ~np.all(masks_RGB == [0, 0, 0], axis=2)
    origin_np = np.array(source).astype(dtype=np.uint8)
    origin_np[masks_] = origin_np[masks_] * (1 - blend) + masks_RGB[masks_] * blend
    result = np.clip(origin_np, 0, 255).astype(np.uint8)
    result = Image.fromarray(result)
    draw_bbox_label(result, boxes, [default_classes[label] for label in labels], scores,
                    show_bbox, show_labels, show_conf)
    # result.save(os.path.join(save_dir, os.path.splitext(filename)[0] + ".png"))
    result = result.resize(origin_size, Image.NEAREST)
    count = Counter([default_classes[label] for label in labels])
    count = {cls: count.get(cls, 0) for cls in default_classes}
    return result, cv2.resize(mask_np, origin_size), count


def seg_predict(source, colors=None,
                model_cfg: Union[str, mmengine.Config] = 'weights/PIDNet/pidnet-s_2xb6-120k_512x512-roadcracks.py',
                deploy_cfg: Union[str, mmengine.Config] = 'weights/PIDNet/segmentation_onnxruntime_static-512x512.py',
                backend_files=None,
                device: str = 'cpu'):
    """

    :param source: 图片 Input image file or numpy array for inference or url/image
    :param colors: [[]]
    :param model_cfg:
    :param deploy_cfg:
    :param backend_files:
    :param device:
    :return:
    """
    if backend_files is None:
        backend_files = ['weights/PIDNet/pidnet-s_2xb6-120k_512x512-roadcracks.onnx']
    download_file('https://file.zouran.top/weights/PIDNet/pidnet-s_2xb6-120k_512x512-roadcracks.onnx',backend_files[0])
    if colors is None:
        colors = [[0, 0, 0], [255, 255, 255]]
    result = PIDNet(source, model_cfg, deploy_cfg, backend_files, device)[0]
    return draw_RGB_mask(result, colors)
